#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#define FILTERDIM 5
#define FILTERSIZE 25
#define TILEWIDTH 16
#define RADIUS 2
#define HOSTS 4

typedef unsigned char uchar;

struct PPMstructure {
  char type[2];
  int width;
  int height;
  int maxVal;
  uchar *r;
  uchar *g;
  uchar *b;
  int *k;
};

typedef struct PPMstructure PPMImage;

/* Verifica erros */
void cudaCheck(cudaError_t error) {
    if(error != cudaSuccess) {
        fprintf(stderr,"ERROR: %s\n", cudaGetErrorString(error));
        exit(1);
    }
}

/* Aloca memoria */
void allocData(PPMImage *img, char* type, int width, int height, int maxVal) {
    strcpy(img->type, type);
    img->width = width;
    img->height = height;
    img->maxVal = maxVal;

    img->r = (uchar *)malloc(width * height * sizeof(uchar));
    img->g = (uchar *)malloc(width * height * sizeof(uchar));
    img->b = (uchar *)malloc(width * height * sizeof(uchar));
    img->k = (int *)malloc(width * height * sizeof(int));
}

/* Libera memoria */
void freeData(PPMImage *img) {
    free(img->r);
    free(img->g);
    free(img->b);
    free(img->k);
}

/* Cria estrutura PPMImage */
void readImage(PPMImage *imgIn, PPMImage *imgOut, PPMImage *vecIn, PPMImage *vecOut, char *in){
    FILE *input;
    int i, j, k, z;
    char type[2];
    int maxVal, width, height, blocWidth, blocHeight;

    input = fopen(in, "r");

    if (input == NULL)
        printf("error\n");
    else {
        /* Le header */
        // Le tipo
        fscanf(input, "%s", type);
        // Pula fim da linha
        while (getc(input) != '\n');

        // Pula comentario
        while (getc(input) == '#') {
            while (getc(input) != '\n');
        }

        // Volta um caracter
        fseek(input, -1, SEEK_CUR);

        // Le tamanho
        fscanf(input, "%d", &width);
        fscanf(input, "%d", &height);
        fscanf(input, "%d", &maxVal);

        allocData(imgIn, type, width, height, maxVal);
        allocData(imgOut, type, width, height, maxVal);

        // Le dados e separa em canais
        for(i = 0; i < width * height ; i++) {
            fscanf(input, "%hhu %hhu %hhu", &(imgIn->r[i]), &(imgIn->g[i]), &(imgIn->b[i]));
        }

        // Recalcula dimensoes das imagens menores
        blocWidth = width / (HOSTS / 2);
        blocHeight = height / (HOSTS / 2);

        // Aloca memoria para as imagens menores
        for(i = 0; i < HOSTS ; i++) {
            allocData(&vecIn[i], type, blocWidth+RADIUS, blocHeight+RADIUS, maxVal);
            allocData(&vecOut[i], type, blocWidth+RADIUS, blocHeight+RADIUS, maxVal);
        }

        // Cria as imagens menores
        z = 0;
        for(i = 0; i < blocHeight + RADIUS; i++) {
            for(j = 0; j < blocWidth + RADIUS; j++) {
                k = (i * width) + j;
                vecIn[0].r[z] = imgIn->r[k];
                vecIn[0].g[z] = imgIn->g[k];
                vecIn[0].b[z] = imgIn->b[k];
                vecIn[0].k[z] = k;
                
                k = (blocWidth - RADIUS) + ((i * width) + j);
                vecIn[1].r[z] = imgIn->r[k];
                vecIn[1].g[z] = imgIn->g[k];
                vecIn[1].b[z] = imgIn->b[k];
                vecIn[1].k[z] = k;
                
                k = ((blocHeight - RADIUS) * width) + ((i * width) + j);
                vecIn[2].r[z] = imgIn->r[k];
                vecIn[2].g[z] = imgIn->g[k];
                vecIn[2].b[z] = imgIn->b[k];
                vecIn[2].k[z] = k;
                
                k = ((blocHeight - RADIUS) * width) + (blocWidth - RADIUS) + ((i * width) + j);
                vecIn[3].r[z] = imgIn->r[k];
                vecIn[3].g[z] = imgIn->g[k];
                vecIn[3].b[z] = imgIn->b[k];
                vecIn[3].k[z] = k;

                z++;
            }
        }
    }
    fclose(input);
}

/* Cria arquivo com o resultado */
void saveImage(PPMImage *img, PPMImage *vec, char *out) {
    FILE *output;
    int i, j;

    // Escreve header
    output = fopen(out, "w");
    fprintf(output, "%s\n", img->type);
    fprintf(output, "%d %d\n", img->width, img->height);
    fprintf(output, "%d\n", img->maxVal);

    // Refaz a imagem grande com as menores
    for (i = 0; i < vec->height; i++) {
        for (j = 0; j < vec->width; j++) {
            img->r[vec[0].k[(i * (vec->width)) + j]] = vec[0].r[(i * (vec->width)) + j];
            img->g[vec[0].k[(i * (vec->width)) + j]] = vec[0].g[(i * (vec->width)) + j];
            img->b[vec[0].k[(i * (vec->width)) + j]] = vec[0].b[(i * (vec->width)) + j];

            img->r[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].r[(i * (vec->width)) + j + RADIUS];
            img->g[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].g[(i * (vec->width)) + j + RADIUS];
            img->b[vec[1].k[(i * (vec->width)) + j + RADIUS]] = vec[1].b[(i * (vec->width)) + j + RADIUS];

            img->r[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].r[((i + RADIUS) * (vec->width)) + j];
            img->g[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].g[((i + RADIUS) * (vec->width)) + j];
            img->b[vec[2].k[((i + RADIUS) * (vec->width)) + j]] = vec[2].b[((i + RADIUS) * (vec->width)) + j];

            img->r[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].r[((i + RADIUS) * (vec->width)) + j + RADIUS];
            img->g[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].g[((i + RADIUS) * (vec->width)) + j + RADIUS];
            img->b[vec[3].k[((i + RADIUS) * (vec->width)) + j + RADIUS]] = vec[3].b[((i + RADIUS) * (vec->width)) + j + RADIUS];
        }
    }

    // Escreve dados
    for (i = 0; i < img->height * img->width; i++) {
        fprintf(output, "%hhu %hhu %hhu ", img->r[i], img->g[i], img->b[i]);
    }

    fclose(output);
}

__global__ void filter(uchar *in, uchar *out, int H, int W) {
    int i, j, k;

    int x = blockDim.x * blockIdx.x + threadIdx.x;
    int y = blockDim.y * blockIdx.y + threadIdx.y;

    int idx = y * W + x;
    float Pvalue = 0;

    k = FILTERDIM / 2;

    // Restringe bordas
    if (x >= W-k|| y >= H-k|| x <= k-1 || y <= k-1) return;

    // Convolucao
    for(i = 0; i <= FILTERDIM - 1; i++){
        for(j = 0; j <= FILTERDIM - 1; j++){
           Pvalue += in[(idx - k + i) + W * (j - k)];
        }
    }
    out[idx] = (uchar)(Pvalue / FILTERSIZE);
}


int main(int argc, char const *argv[]) {
    PPMImage imgIn, imgOut, *vecIn, *vecOut;
    uchar *inR[HOSTS], *inG[HOSTS], *inB[HOSTS], *outR[HOSTS], *outG[HOSTS], *outB[HOSTS];
    char in[20], out[20];
    int i;

    strcpy(in, argv[1]);
    strcpy(out, argv[2]);

    vecIn = (PPMImage *)malloc(HOSTS * sizeof(PPMImage));
    vecOut = (PPMImage *)malloc(HOSTS * sizeof(PPMImage));

    readImage(&imgIn, &imgOut, vecIn, vecOut, in);

    int size = vecIn[0].width * vecIn[0].height;

    for (i = 0; i < HOSTS; i++) {
        cudaCheck(cudaMalloc((void**)&inR[i], size * sizeof(uchar)));
        cudaCheck(cudaMalloc((void**)&inG[i], size * sizeof(uchar)));
        cudaCheck(cudaMalloc((void**)&inB[i], size * sizeof(uchar)));
        cudaCheck(cudaMalloc((void**)&outR[i], size * sizeof(uchar)));
        cudaCheck(cudaMalloc((void**)&outG[i], size * sizeof(uchar)));
        cudaCheck(cudaMalloc((void**)&outB[i], size * sizeof(uchar)));

        cudaCheck(cudaMemcpy(inR[i], vecIn[i].r, size * sizeof(uchar), cudaMemcpyHostToDevice));
        cudaCheck(cudaMemcpy(inG[i], vecIn[i].g, size * sizeof(uchar), cudaMemcpyHostToDevice));
        cudaCheck(cudaMemcpy(inB[i], vecIn[i].b, size * sizeof(uchar), cudaMemcpyHostToDevice));

        dim3 gridDim(vecIn[i].width / TILEWIDTH + 1, vecIn[i].height / TILEWIDTH + 1);
        dim3 blockDim(TILEWIDTH, TILEWIDTH);

        filter<<<gridDim, blockDim>>>(inR[i], outR[i], vecIn[i].height, vecIn[i].width);
        filter<<<gridDim, blockDim>>>(inG[i], outG[i], vecIn[i].height, vecIn[i].width);
        filter<<<gridDim, blockDim>>>(inB[i], outB[i], vecIn[i].height, vecIn[i].width);

        cudaCheck(cudaMemcpy(vecOut[i].r, outR[i], size * sizeof(uchar), cudaMemcpyDeviceToHost));
        cudaCheck(cudaMemcpy(vecOut[i].g, outG[i], size * sizeof(uchar), cudaMemcpyDeviceToHost));
        cudaCheck(cudaMemcpy(vecOut[i].b, outB[i], size * sizeof(uchar), cudaMemcpyDeviceToHost));
    }

    saveImage(&imgOut, vecOut, out);

    for (i = 0; i < HOSTS; i++) {
        cudaFree(inR[i]);
        cudaFree(inB[i]);
        cudaFree(inG[i]);
        cudaFree(outR[i]);
        cudaFree(outG[i]);
        cudaFree(outB[i]);
    }

    freeData(&imgIn);
    freeData(&imgOut); 

    free(vecIn);
    free(vecOut);
    
    return 0;
}

